package addressBook;

import java.util.Iterator;
import java.util.LinkedList;
import java.util.NoSuchElementException;
import java.util.Stack;
import java.util.function.Consumer;


/**
 * 二叉树的迭代器.
 */
public class TreeIterator<E> implements Iterator<E> {
    /**
     * 遍历的根节点
     */
    private TreeNode<E> root;
    /**
     * 储存输出的结果
     */
    public LinkedList<TreeNode<E>> list;

    /**
     * 自定义比遍历方式，默认是中序遍历
     * 1 为 preOrder,2 为 inOrder,3 为 postOrder
     */
    private int traverseMethod;
    /**
     * 配合链表list的指针，用来判断list长度和list是否为空.
     */
    int index;


    public TreeIterator(TreeNode<E> root) {
        this(root, 2);
    }

    public TreeIterator(TreeNode<E> root, int traverseMethod) {
        list = new LinkedList<>();
        this.root = root;
        this.traverseMethod = traverseMethod;
        this.index = -1;
        traverse();

    }

    /**
     * 遍历二叉树并储存到list里面.
     */
    private void traverse() {
        if (traverseMethod == 1) {
            preOrder(this.root);
        } else if (traverseMethod == 3) {
            postOrder(this.root);
        } else {
            inOrder(this.root);
        }
    }

    /**
     * 前序遍历
     *
     * @param root 根节点root
     */
    private void preOrder(TreeNode<E> root) {
        //创建一个栈
        Stack<TreeNode<E>> stack = new Stack<>();
        //遍历到 最后一个节点的时候 ，该节点的左右子树都为空，且栈中的值应该也为空，
        //所以只要不是同时满足这2点都要进入循环。
        while (root != null || !stack.isEmpty()) {
            //遍历的顺序得知，先需要一直往左走，
            while (root != null) {
                //输出语句
                list.add(root);
                //压栈
                stack.push(root);
                //遍历左子树
                root = root.left;
            }
            //遍历顺序得知：左子树遍历完成，弹出父节点，然后遍历父节点的右子树.
            if (!stack.isEmpty()) {
                //左子树遍历结束后弹出左子树的父节点
                root = stack.pop();
                //遍历右子树
                root = root.right;
            }
        }
    }

    /**
     * 中序遍历，输出一个有序的集合。
     *
     * @param root 根节点
     */
    private void inOrder(TreeNode<E> root) {
        Stack<TreeNode<E>> stack = new Stack<>();
        while (root != null || !stack.isEmpty()) {
            while (root != null) {
                stack.push(root);
                //遍历左子树
                root = root.left;
            }
            if (!stack.isEmpty()) {
                TreeNode<E> popNode = stack.pop();
                //输出语句
                list.add(popNode);
                //遍历右子树
                root = popNode.right;
            }
        }
    }


    /**
     * 后序遍历
     *
     * @param root 根节点root
     */
    private void postOrder(TreeNode<E> root) {
        Stack<TreeNode<E>> stack = new Stack<>();
        //判断是否已经遍历过的指针指向已经遍历过的节点.
        TreeNode<E> visited = root;
        while (root != null || !stack.isEmpty()) {
            while (root != null) {
                //储存遍历的节点在栈stack中.
                stack.push(root);
                //总是优先遍历左子树.
                root = root.left;
            }
            //peek节点进行判断，而不是直接弹出，这里很重要。
            root = stack.peek();
            //if node(root).right is null or the right is visited
            // (we  out put it and let the visited = node)
            // and then ,we pop the stack top to invoke the up level node.
            // 如果节点为空或者节点已经visited都可以输出
            // 输出之后我们让visited等于已经输出的子树
            if (root.right == null || root.right == visited) {
                list.add(root);
                visited = root;
                root = stack.pop();
                //这里防止死循环.
                root = null;
            } else {
                //如果右子树不为空并且没有visited就遍历右子树.
                root = root.right;
            }
        }
    }


    @Override
    public boolean hasNext() {
        return !list.isEmpty() && index < (list.size() - 1);
    }


    @Override
    public E next() {
        index++;
        return list.get(index).element;
    }


    @Override
    public void remove() {
        list.remove(index);
    }


}
